package main

import (
	"flag"
	"fmt"
	"os"
	"network"
	"strconv"
	"strings"
	"time"
	"sync"
)

//输入参数列表：
//-n : 并发连接数，默认1
//-h : 连接主机IP/域名，必填，支持https，如：https://baidu.com
//-p : 端口号，默认80
//-url : 地址，默认/，如：/info/index
//-m : 请求类型，默认GET，GET,POST等
//-args : 请求参数，key=value格式，多个用&分隔 或者 json格式（请自行URL编码）
//-argt : 请求参数类型，json/form
//-attd : 附加数据包大小（单位:KB），主要用于网络带宽测试
//-showr : 是否打印响应结果，0=不打印，1=打印响应状态，2=打印全部响应内容

func main() {

	//启动参数定义
	n := flag.Int("n", 1, "并发连接数")
	h := flag.String("h", "", "连接主机IP/域名")
	p := flag.Int("p", 80, "端口号，默认80")
	url := flag.String("url", "/", "URL")
	m := flag.String("m", "GET", "请求方法类型")
	args := flag.String("args", "", "请求参数")
	argt := flag.String("argt", "", "请求参数类型")
	attd := flag.Int("attd", 0, "附加数据包大小（单位:KB） >0")
	showr := flag.Int("showr", 0, "是否打印响应结果，0=不打印，1=打印响应状态，2=打印全部响应内容")

	//参数解析
	flag.Parse()

	//参数打印
	fmt.Println("并发连接数：", *n)
	fmt.Println("连接主机IP/域名：", *h)
	fmt.Println("端口号：", *p)
	fmt.Println("URL：", *url)
	fmt.Println("请求方法类型：", *m)
	fmt.Println("请求参数：", *args)
	fmt.Println("请求参数类型：", *argt)
	fmt.Println("附加数据包大小：", *attd)
	fmt.Println("是否打印响应结果：", *showr)

	//检查参数
	parmaeCheck(*n, *p, *h, *url, *m, *args, *argt, *attd, *showr)

	//创建网络连接
	fmt.Println("开始集中创建网络连接....")
	createNetConn(*n, *p, *h)
	fmt.Println("网络连接集中创建完成！")

	var wg sync.WaitGroup
	wg.Add(1)

	//启动HTTP请求
	fmt.Println("开始并发请求数据....")
	startHttpReq(*h, *p, *url, *m, *args, *argt, *attd, *showr)

	wg.Wait()
}


//参数检查
func parmaeCheck(num ,port int, host, url, method, args, argt string, attd, showr int) {
	var methods = make(map[string]string)
	methods["GET"] = "GET"
	methods["POST"] = "POST"
	methods["HEAD"] = "HEAD"
	methods["OPTIONS"] = "OPTIONS"
	methods["PUT"] = "PUT"
	methods["PATCH"] = "PATCH"
	methods["DELETE"] = "DELETE"
	methods["TRACE"] = "TRACE"
	methods["CONNECT"] = "CONNECT"

	var isErr bool = false
	if num < 1 || num > 1000000 {
		fmt.Println("并发连接数在1~1000000之间.")
		isErr = true
	}
	if host == "" {
		fmt.Println("连接主机IP/域名必须输入.")
		isErr = true
	}
	if port < 1 || port > 65535 {
		fmt.Println("端口号输入不正确.")
		isErr = true
	}
	if url == "" {
		fmt.Println("URL输入不正确.")
		isErr = true
	}
	if method == "" {
		fmt.Println("请求方法类型输入不正确.")
		isErr = true
	} else if _,OK := methods[method]; !OK  {
		fmt.Println("请求方法类型输入不正确.")
		isErr = true
	}
	if args != "" {
		if argt != "json" && argt != "form" {
			fmt.Println("请求参数类型输入不正确.")
			isErr = true
		}
	}
	if attd < 0 || attd > 1024000 {
		fmt.Println("附加数据包在0~102400之间.")
		isErr = true
	}
	if showr != 0 && showr != 1 && showr !=2 {
		fmt.Println("打印结果标识输入不正确：0=不打印，1=打印响应状态，2=打印全部响应内容.")
		isErr = true
	}
	if isErr {
		os.Exit(0)
	}
}


//客户端网络列表
var connList []*network.TcpClient
//创建网络连接
func createNetConn(num ,port int, host string) {
	var okConnNum = 0
	var failConnNum = 0
	for i:=0 ; i<num;i++ {
		conn := &network.TcpClient{}
		conn.Name = fmt.Sprintf("连接%d号", i)
		conn.Host = host
		conn.Port = port
		OK := conn.Init()
		if OK {
			okConnNum += 1
		} else {
			failConnNum += 1
		}
		connList = append(connList, conn)
	}
	fmt.Println(fmt.Sprintf("创建网络连接成功[%d],失败[%d],总共[%d].",okConnNum,failConnNum,num))
}


//请求数据包对象
var httpReqData *network.HttpReqData
var httpReqContent string
var httpReqContentByte []byte
//启动HTTP请求
func startHttpReq(host string, port int, url, method, args, argt string, attd, showr int) {
	//组装请求数据
	httpReqData = &network.HttpReqData{}
	httpReqData.Method = method
	httpReqData.URL = url
	httpReqData.Version = "HTTP/1.1"
	httpReqData.Accept = "*/*"
	host = strings.Replace(host,"https://","",1)
	host = strings.Replace(host,"http://","",1)
	httpReqData.Host = host + ":" + strconv.Itoa(port)
	httpReqData.Cache_Control = "no-cache"
	httpReqData.Connection = "Keep-Alive"
	httpReqData.Refer = httpReqData.Host
	httpReqData.User_Agent = "Mozilla/5.0(Macintosh;IntelMacOSX10_7_0)AppleWebKit/535.11(KHTML,likeGecko)Chrome/17.0.963.56Safari/535.11"
	if argt == "json" {
		httpReqData.Content_Type = "application/json"
	} else if argt == "form" {
		httpReqData.Content_Type = "application/x-www-form-urlencoded"
	}
	//添加扩展数据包
	if attd > 0 {
		ad := network.GetAttd(attd)
		if argt == "json" {
			newstr := fmt.Sprintf("{\"_attd_\":\"%s\",",ad)
			args = strings.Replace(args,"{", newstr,1)
		} else if argt == "form" {
			httpReqData.Content_Type = "application/x-www-form-urlencoded"
			args += "&_attd_=" + ad
		} else {
			args += "_attd_=" + ad
		}
	}
	//计算正文长度
	content_leng := len(args)
	httpReqData.Content_Length = strconv.Itoa(content_leng)
	httpReqData.Content = args

	//打印请求数据包
	httpReqContent = httpReqData.ToString()
	fmt.Println("请求数据包上下文：")
	fmt.Println(httpReqContent)

	//创建携程，并发请求
	httpReqContentByte = []byte(httpReqContent)
	var connleng = len(connList)
	for i:=0;i<connleng;i++{
		channel := make(chan *network.TcpClient,3)
		conn := connList[i]
		channel <- conn
		go func(ch chan *network.TcpClient) {
			myconn := <- ch
			if myconn == nil {
				fmt.Println("tcpConn对象为空.")
				return
			}
			for {
				defer func() {
					p := recover()
					if p != nil {
						fmt.Println(fmt.Sprintf("%s 循环发生请求时发生异常：", myconn.Name), p)
					}
				}()

				myconn.SendHttpReq(httpReqContentByte, showr)
				//休眠10秒
				time.Sleep(time.Second * 10)
			}
		}(channel)

	}
}


